from flask import Flask, request, jsonify, send_from_directory
from threading import Thread
import time
import os
import uuid
from scripts.my_sample import load_pipe, infer
from constant import port
# Initialize the model pipeline
load_pipe()
print('模型加载完毕！')

app = Flask(__name__)

# Dictionary to store task information
tasks = {}

# Directory to save generated images
output_base_dir = 'output_images'
if not os.path.exists(output_base_dir):
    os.makedirs(output_base_dir)

def run_inference(task_id, prompt, width, height, num_inference_steps, guidance_scale, num_images_per_prompt, seed):
    tasks[task_id]['start_time'] = time.time()
    output_dir, images = infer(
        prompt,
        width,
        height,
        num_inference_steps,
        guidance_scale,
        num_images_per_prompt,
        seed,
    )
    tasks[task_id]['output_dir'] = output_base_dir
    for i in range(len(images)):
        img = images[i]
        img.save(os.path.join(output_base_dir,f'{task_id}_{i}.jpg'))
    tasks[task_id]['output_path'] = [os.path.join(output_base_dir,f'{task_id}_{i}.jpg') for i in range(num_images_per_prompt)]
    tasks[task_id]['status'] = 'completed'
    tasks[task_id]['end_time'] = time.time()
    tasks[task_id]['duration'] = round(tasks[task_id]['end_time'] -  tasks[task_id]['start_time'],2) 

@app.route('/start_inference', methods=['POST'])
def start_inference():
    data = request.json
    prompt = data.get('prompt')
    width = data.get('width', 1024)
    height = data.get('height', 1024)
    num_inference_steps = data.get('num_inference_steps', 32)
    guidance_scale = data.get('guidance_scale', 5.0)
    num_images_per_prompt = data.get('num_images_per_prompt', 1)
    seed = data.get('seed', 90)
    

    if not prompt:
        return jsonify({'error': 'Prompt is required'}), 400

    task_id = str(uuid.uuid4())
    task_id = str(data.get('task_id', task_id))
    tasks[task_id] = {'status': 'running'}

    thread = Thread(target=run_inference, args=(task_id, prompt, width, height, num_inference_steps, guidance_scale, num_images_per_prompt, seed))
    thread.start()

    return jsonify({'task_id': task_id})
@app.route('/get_image_info/<task_id>', methods=['GET'])
def get_image_info(task_id):
    task_id = str(task_id)
    task = tasks.get(task_id)
    if not task:
        return jsonify({'error': 'Task ID not found'}), 404
 
    return jsonify(task), 200
@app.route('/get_image/<task_id>/<image_i>', methods=['GET'])
def get_image(task_id,image_i):
    task_id = str(task_id)
    task = tasks.get(task_id)
    if not task:
        return jsonify({'error': 'Task ID not found'}), 404

    if task['status'] != 'completed':
        return jsonify({'status': task['status']}), 202

    output_dir = task['output_dir']
    image_files = os.listdir(output_dir)
    #if not image_files:
    #    return jsonify({'error': 'No images found'}), 404
    #
    image_path = os.path.join(output_dir, f'{task_id}_{image_i}.jpg')
    if  not os.path.exists(image_path):
        return jsonify({'error': 'No images found'}), 404
    return send_from_directory(output_dir, f'{task_id}_{image_i}.jpg')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=port)
